import networkx as nx
import collections
import functools

class memoize:
  # From https://github.com/S2E/s2e-env/blob/master/s2e_env/utils/memoize.py

  def __init__(self, func):
    self._func = func
    self._cache = {}

  def __call__(self, *args):
    if not isinstance(args, collections.abc.Hashable):
      return self._func(args)

    if args in self._cache:
      return self._cache[args]

    value = self._func(*args)
    self._cache[args] = value
    return value

  def __repr__(self):
    # Return the function's docstring
    return self._func.__doc__

  def __get__(self, obj, objtype):
    # Support instance methods
    return functools.partial(self.__call__, obj)

dotfilename = './callgraph.dot'

G = nx.DiGraph(nx.drawing.nx_pydot.read_dot(dotfilename))
print(nx.info(G))

def node_name(name):
    return "\"{%s}\""%name

@memoize
def find_nodes(name):
    n_name = node_name(name)
    return [(n,d) for n,d in G.nodes(data=True) if n_name in d.get('label', '')]

main = find_nodes("main")
main_callsite_set = set()
for n,d in main:
    print(n, G.nodes[n])
    # res = G.successors(n)
    res = nx.bfs_successors(G, n,depth_limit=1)
    for node,succ in res:
        # print(G.nodes[succ]['label'])
        # main_callsite_set.add(G.nodes[succ]['label'])        
        print(G.nodes[node]['label'])
        main_callsite_set.add(G.nodes[node]['label'])
    print('↑'*30) 
    
    res = nx.bfs_predecessors(G, n, depth_limit=1)
    for node, pred in res:
        print(G.nodes[node]['label'])
    
    print('-'*30)   
    cur = G.nodes[n]['label']
    print('predecessor')
    for node in G.predecessors(n):
        print('{}->{}'.format(cur, G.nodes[node]['label']))
        
    print('-'*30)    
    print('successors')
    for node in G.successors(n):
        print('{}->{}'.format(cur, G.nodes[node]['label']))
        
print('-'*30)    

res = find_nodes('xmlBufAdd')
for n,d in res:
    print(n)
    print(G.nodes[n])
    res = G.predecessors(n)
    
    for node,pred in nx.bfs_predecessors(G, n, depth_limit=1):
        print("{}<-{}".format(G.nodes[node]['label'], G.nodes[pred]['label']))
        # if G.nodes[node]['label'] in main_callsite_set:
        #     print('___^^^___')
    print('-'*30)
    for node,succes in nx.bfs_successors(G, n, depth_limit=1):
        print(G.nodes[node]['label'])
        cur = G.nodes[node]['label']
        for succ in succes:
            print("{}->{}".format(cur, G.nodes[succ]['label']))
        # if G.nodes[node]['label'] in main_callsite_set:
        #     print('___^^^___')
    print('-'*30)   
    cur = G.nodes[n]['label']
    for node in G.predecessors(n):
        print('{}->{}'.format(cur, G.nodes[node]['label']))
        
    print('-'*30)    
    for node in G.successors(n):
        print('{}->{}'.format(cur, G.nodes[node]['label']))

print('finish search')